#ifndef RKNN_ENGINE_H
#define RKNN_ENGINE_H

#include "rknn_api.h"

// --- Model-specific Parameters (Updated for YOLOv8s) ---
#define MODEL_INPUT_WIDTH 640
#define MODEL_INPUT_HEIGHT 640
#define NMS_THRESHOLD 0.40      // Corresponds to your Python script
#define CONF_THRESHOLD 0.55     // Corresponds to your Python script
#define NUM_CLASSES 3
#define MAX_DETECT_RESULTS 256
#define MAX_RKNN_OUTPUT_COUNT 10


// --- Data Structures ---

/**
 * @brief Stores a single detection result.
 */
typedef struct
{
    int left;
    int top;
    int right;
    int bottom;
    float confidence;
    int class_id;
    const char *class_name;
} detect_result_t;

/**
 * @brief Stores a group of detection results for a single frame.
 */
typedef struct
{
    int count;
    detect_result_t results[MAX_DETECT_RESULTS];
} detect_result_group_t;

/**
 * @brief Context for the RKNN inference engine.
 *
 * Holds all resources related to the NPU, including the RKNN context,
 * model data, and tensor attributes.
 */
typedef struct
{
    rknn_context ctx;
    unsigned char *model_data;
    int model_data_size;

    rknn_input_output_num io_num;
    rknn_tensor_attr *input_attrs;
    rknn_tensor_attr *output_attrs;

    rknn_tensor_mem *output_mems[MAX_RKNN_OUTPUT_COUNT];
    int model_channel;
    int model_width;
    int model_height;

} rknn_engine_context_t;


// --- Public API Functions ---

/**
 * @brief Initializes the RKNN model and context.
 * @param eng_ctx Pointer to the engine context structure to be initialized.
 * @param model_path Path to the .rknn model file.
 * @return 0 on success, -1 on failure.
 */
int rknn_engine_init(rknn_engine_context_t *eng_ctx, const char *model_path);

/**
 * @brief Runs inference on a single input frame.
 * @param eng_ctx Pointer to the initialized engine context.
 * @param input_mem Pointer to the rknn_tensor_mem structure describing the input buffer (DMABUF).
 * @param results Pointer to a structure where the detection results will be stored.
 * @param frame_width The original width of the video frame (e.g., 1920).
 * @param frame_height The original height of the video frame (e.g., 1080).
 * @return 0 on success, -1 on failure.
 */
int rknn_engine_run(rknn_engine_context_t *eng_ctx, rknn_tensor_mem *input_mem, detect_result_group_t *results,
                    int frame_width, int frame_height);

/**
 * @brief Releases all resources associated with the RKNN engine.
 * @param eng_ctx Pointer to the engine context to be released.
 */
void rknn_engine_release(rknn_engine_context_t *eng_ctx);

#endif // RKNN_ENGINE_H
